#!/usr/bin/env python
# -*- coding: UTF-8 -*-

'''
Copyright (C) 2010 dbzhang800@gmail.com
All rights reserved.

2010.07.21-2010.07.22
''' 

import shutil
import os
import os.path
import sys
import logging
import time
from msexcel import *
from aber import *
from soelens import *
from textcolor import *

class DemoConfig:
    u"""
     * 配置文件说明：这是个配置文件的demo，你可以参考 real_config_example.py 编写一个 real_config.py 文件，放于当期目录即可
    1. munroPath         : 指定 munro 程序在所目录(可以是字符串，也可以是一个列表)
    2. workPath          : 指定 要计算的数据所在目录，使用相对路径时，相对本脚本所在位置
    3. aberDatFileName   : 指定 aber 的.dat的文件名，可省略后缀
    4. xlsxFileName      : 指定 xlsx 文件文件路径，使用相对路径时，相对脚本所在位置
    5. elensFilePerfix   : 静电透镜文件名的前缀
    6. realApertureDiameter: 光阑大小列表
     * 其他选项：(调试用，或其他需求)
    7. create_sub_folder : 是否创建文件夹(每一种情况放置一个文件夹，还是不创建任何文件夹)
    8. del_depands       : 若创建文件夹(计算时，需要的 .axb、.axv等都从上层文件夹会拷进来，控制计算结束后删除与否)
    9. pause_calc        : 调用一次 aber.exe 后暂停，按回车键继续
    a. check_axv_files   : .axv 文件是手动生成的，程序中是直接根据规则生成的，确保二者对应(axv或其对应的dat存在一个即可)
    b. skip_if_res_exists: 如果结果文件已经存在，跳过这次，直接进行下次计算
    c. is_demo           : 不对aber的dat文件进行修改(调试或演示使用)
    """
    munroPath = [r"E:\opee pack-0.3.0.1\data\munro" , u"F:\\整理\\ziliao\\linux\\work\\munro\\munro", r"E:\munro"]
    workPath  = "root"
    aberDatFileName   = "aberd.dat"
    xlsxFileName  = u"WD与物镜激励及FOV视场关系.xlsx"
    elensFilePerfix = "soelens1"
    realApertureDiameter = [0.001, 0.002, 0.003, 0.005, 0.007, 0.01, 0.015, 0.02, 0.025,
            0.03, 0.04, 0.05, 0.07, 0.1, 0.12, 0.15, 0.20, 0.25]
    create_sub_folder = True
    del_depands = False
    pause_calc = False
    check_axv_files = False
    skip_if_res_exists = True
    is_demo = True

try:
    from real_config import Config
except ImportError, err:
    print ("real_config not load, use demo_config")
    raw_input("press enter to continue")
    Config = DemoConfig

def change_aber_dat(elens, le, aberdat):
    u""" 修改aber dat 文件
    1. 修改 静电透镜(index 用来确定是要修改第几个静电透镜的数据，从1开始数)
       * 各个电极的电压，根据 落点能量推算(程序具有通用性，对电极数目不做任何限制)
       * 文件名，根据下面规则生成
    2. 修改 Image Plane， 根据 elens 确定(该功能以从该函数内移除，改为从axv对应的dat内读取)
    """
    if Config.is_demo:
        return
    axv_name = generate_axv_name(elens)
    v1 = le+8000
    v2 = le
    v3 = le

    aberdat.setBeamVoltage(v1)
    aberdat.setELensVoltsAndName((v1, v2, v3), axv_name, index=1) #note
    pass

def generate_axv_name(distance, suffix='axv'):
    u""" 生成 .axv 文件名 "前缀+工作距离*10.axv"
    distance：工作距离
    处理时，先乘以10，如不足3位，前面补0
    """
    if suffix[0] == '.':
        suffix = suffix[1:]
    tmp = str(int(distance*10))
    if len(tmp) == 1:
        tmp = '00'+tmp
    elif len(tmp) == 2:
        tmp = '0'+tmp
    return "%s-%s.%s" % (Config.elensFilePerfix, tmp, suffix)

def dat_from_aber_res(resfile):
    u"""
    从 aber的计算结果 .res 获取所需要的两个数据
    1. 激励值：通过 index 指定是读取第几个磁透镜的激励值
    2. 灵敏度：通过 field 指定是读取main、还是sub 的灵敏度

    dat1, dat2 是我们要写入到 .xlsx 文件的两个值
    """
    if isinstance(resfile, str):
        resfile = AberResFile(resfile)
    
    excitation = resfile.getMLensExcitation(index = 1)
    sensitivities = resfile.getDeflectionSensitivities(field='main')
    beamspot = resfile.getBeamSpotDiameter()

    dat1 = excitation
    dat2 = sensitivities
    dat3 = beamspot

    return dat1, dat2, dat3 

def one_aber_calc(elens, le, aperture ,aberdat):
    u"""  进行一次 aber 的计算
    elens:   静电透镜文件名(或能合成需要的文件名的字符串) 
    le:      落点能量信息
    aberdat: aber 的 dat 文件的对象
    """
    # aberdat :save as new file
    aberdat_new_name = generate_aber_new_name(elens, le, aperture)
    aberdat.save(aberdat_new_name)
    # copy aber dat depend files
    if Config.create_sub_folder:
        depends = aberdat.getDependFileNames()
        for d in depends:
            try:
                shutil.copyfile(os.path.abspath("../../../%s" % d), d)
            except:
                shutil.copyfile(os.path.abspath("../../%s" % d), d)
    # call aber.exe
    callMunroAber(aberdat_new_name)
    # del aber dat depend files
    if Config.create_sub_folder and Config.del_depands:
        for d in depends:
            try:
                os.remove(d)
            except:
                pass

def one_soelens_calc(elens):
    u""" 进行一次 solens 的计算
    elens:      静电透镜文件名(或能合成需要的文件名的字符串) 
    le:         落点能量信息
    soelensdat: soelensdat 的 dat 文件的对象
    """
    f = generate_axv_name(elens, 'dat')
    print_info(f)
    callMunroSoelens(f)
    if Config.create_sub_folder and Config.del_depands:
        try:
            os.remove(f)
        except:
            pass

def info_from_xlsx(xlsfile):
    u""" 从 xlsx 中读取所需要的数据
    1. 电透镜文件的信息(列表)
    2. 落点能量数据的列表
    """
    elens_list = xlsfile.row(1)
    elens_list.pop(0)
    le_list = xlsfile.col(1)
    le_list.pop(0)

    return elens_list, le_list

def write_data_to_xlsx(v1, v2, row, col, xlsfile):
    u""" 将两个值分别写入xlsx的两个Sheets中"""
    xlsfile.setcell(row, col, v1, 1)
    xlsfile.setcell(row, col, v2, 2)
    xlsfile.save()

def generate_aber_new_name(elens, le, aperture, suffix=None):
    name =  "aber_%s_%s_%s" % (elens, le, aperture)
    if suffix:
        if suffix[0]!='.':
            suffix = '.'+suffix
        name += suffix
    return name

def print_error(str):
    DCC.setColor(DCC.red)
    print str
    DCC.resetColor()

def print_info(str):
    DCC.setColor(DCC.green)
    print str
    DCC.resetColor()

def gen_le_dir_name(le):
    return "le"+str(int(le))

def gen_wd_dir_name(elens):
    return "distance"+str(elens)

def gen_aperture_d_dir_name(aperture):
    return "aperture"+str(aperture)

def gen_result_file_name(elens, le):
    return  "res_%s_%s.txt" % (elens, le)

if __name__ == '__main__':
    setMunroPath(Config.munroPath)
    xlsfile = MSExcel(Config.xlsxFileName)
    elens_list, le_list = info_from_xlsx(xlsfile)
    if Config.is_demo:
        del xlsfile

    workpath = os.path.join(os.path.curdir, Config.workPath) 
    os.chdir(workpath)

    logging.basicConfig(level = logging.DEBUG,
            format =' %(asctime)s %(levelname)s %(message)s',
            filename = time.strftime('%Y%m%d%H%M%S') + '.log',
            filemode = 'w')

    logging.info("start check axv files")


    if Config.check_axv_files:
        logging.info("start check axv files")
        for d in elens_list:
            f_axv = generate_axv_name(d)
            f_dat = generate_axv_name(d, "dat")
            if (not os.path.exists(f_axv)) and (not os.path.exists(f_dat)):
                logging.error("%s not found" % f_dat)
                print_error("Error: %s not found" % f_dat)
                raw_input("Press <Enter> to exit...")
                exit()

    print_info("HPI: OK! axv files check pass")
    logging.info("axv files check pass")

    aberdat = AberDatFile(Config.aberDatFileName)

    for elens_index in range(len(elens_list)):
        elens = elens_list[elens_index]
        #elens_dir = "distance"+str(elens)                 # dir name
        elens_dir = gen_wd_dir_name(elens)
        logging.info("Start deal with %s" % elens)
        print_info ("HPI: OK! Start deal with %s" % elens)
        f_axv = generate_axv_name(elens)
        f_dat = generate_axv_name(elens, 'dat')
        if not Config.is_demo:
            # set image plane
            aberdat.setImagePlane(SoelensDatFile(f_dat).getMeshTableRect().right())
        if Config.create_sub_folder:
            if not os.path.exists(elens_dir):
                os.mkdir(elens_dir)
            os.chdir(elens_dir)
            # if .axv not exists in current dir and parent dir, copy .dat from parent dir
            print (not os.path.exists(f_axv))
            print (not os.path.exists(os.path.abspath("../%s" % f_axv)))
            if (not os.path.exists(f_axv)) and (not os.path.exists(os.path.abspath("../%s" % f_axv))) and (not Config.is_demo):
                shutil.copyfile(os.path.abspath("../%s" % f_dat), f_dat)
                one_soelens_calc(elens)

        for le_index in range(len(le_list)):
            le  = le_list[le_index]
            if Config.create_sub_folder:
                #le_dir = "le"+str(int(le))              # dir name
                le_dir = gen_le_dir_name(le)
                if not os.path.exists(le_dir):
                    os.mkdir(le_dir)
                os.chdir(le_dir) #bug 20100723
            print_info ("HPI: OK! Start deal with %s %s" % (elens, le))
            logging.info("Start deal with %s %s" % (elens, le))

            change_aber_dat(elens, le, aberdat)

            resultlist = []
            for aperture in Config.realApertureDiameter:
                if Config.create_sub_folder:
                    #aperture_dir = "aperture"+str(aperture)              # dir name
                    aperture_dir = gen_aperture_d_dir_name(aperture)
                    if not os.path.exists(aperture_dir):
                        os.mkdir(aperture_dir)
                    os.chdir(aperture_dir)
                if Config.skip_if_res_exists and os.path.exists(generate_aber_new_name(elens, le, aperture, '.res')):
                    logging.info("skip %s %s %s" % (elens, le, aperture))
                    print_info( "HPI: skip %s %s %s" % (elens, le, aperture))
                else:
                    print_info( "HPI: OK! Start deal with %s %s %s" % (elens, le, aperture))
                    logging.info("Start deal with %s %s %s" % (elens, le, aperture))
                    aberdat.setRealApertureDiameter(aperture)
                    one_aber_calc(elens, le, aperture, aberdat)
                v1, v2, v3 = dat_from_aber_res(generate_aber_new_name(elens, le, aperture))
                resultlist.append((aperture, v1, v2, v3));
                if Config.create_sub_folder:
                    os.chdir("..") #aperture
                if Config.pause_calc:
                    raw_input("press <Enter> to continue ...")
            if not Config.is_demo:
                write_data_to_xlsx(resultlist[0][1], resultlist[0][2], le_index+2, elens_index+2, xlsfile)
                logging.info("Write data to xlsx file")
                print_info("Write data to xlsx file")
            #resultfile = "res_%s_%s.txt" % (elens, le)
            resultfile = gen_result_file_name(elens, le)
            with open(resultfile, 'w') as f:
                for vals in resultlist:
                    f.write("%10s %15s %15s %15s\n" % (vals[0], vals[1], vals[2], vals[3]))
            logging.info("Write data to %s" % resultfile)
            print_info("write data to %s" % resultfile)

            if Config.create_sub_folder:
                os.chdir("..") #le
        if Config.create_sub_folder:
            os.chdir("..") #distance

    if not Config.is_demo:
        del xlsfile
